package com.liao.handle;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;

public class ChatWebSocketHandler implements WebSocketHandler {

    private final Set<WebSocketSession> sessions = Collections.synchronizedSet(new HashSet<>());

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.add(session);
        sendMessageToAll(new ChatMessage("SYSTEM", session.getId() + " joined the chat"));
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        ChatMessage chatMessage = new ObjectMapper().readValue(message.getPayload().toString(), ChatMessage.class);
        sendMessageToAll(chatMessage);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        sessions.remove(session);
        sendMessageToAll(new ChatMessage("SYSTEM", session.getId() + " left the chat"));
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        sessions.remove(session);
        sendMessageToAll(new ChatMessage("SYSTEM", session.getId() + " left the chat"));
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    private void sendMessageToAll(ChatMessage chatMessage) throws IOException {
        for (WebSocketSession session : sessions) {
            session.sendMessage(new TextMessage(new ObjectMapper().writeValueAsString(chatMessage)));
        }
    }
}